import os
import argparse
parser = argparse.ArgumentParser()
parser.add_argument('-c', '--cuda', type=str, default='2', help='CUDA device id')
args, unknown = parser.parse_known_args()
os.environ['CUDA_VISIBLE_DEVICES'] = args.cuda

import gradio as gr
import torch
from transformers import AutoConfig, AutoModelForCausalLM
from janus.models import MultiModalityCausalLM, VLChatProcessor
from janus.utils.io import load_pil_images
from PIL import Image

import numpy as np
import os
import time
# import spaces  # Import spaces for ZeroGPU compatibility


# Load model and processor
# model_path = "deepseek-ai/Janus-Pro-7B"

config = None
language_config = None
vl_gpt = None
vl_chat_processor = None



@torch.inference_mode()
def multimodal_understanding(image, question, seed=42, top_p=0.95, temperature=0.1):
    # Clear CUDA cache before generating
    torch.cuda.empty_cache()
    
    # set seed
    torch.manual_seed(seed)
    np.random.seed(seed)
    torch.cuda.manual_seed(seed)
    
    conversation = [
        {
            "role": "<|User|>",
            "content": f"<image_placeholder>\n{question}",
            "images": [image],
        },
        {"role": "<|Assistant|>", "content": ""},
    ]
    
    pil_images = [Image.fromarray(image)]
    prepare_inputs = vl_chat_processor(
        conversations=conversation, images=pil_images, force_batchify=True
    ).to(cuda_device, dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float16)
    
    
    inputs_embeds = vl_gpt.prepare_inputs_embeds(**prepare_inputs)
    
    outputs = vl_gpt.language_model.generate(
        inputs_embeds=inputs_embeds,
        attention_mask=prepare_inputs.attention_mask,
        pad_token_id=tokenizer.eos_token_id,
        bos_token_id=tokenizer.bos_token_id,
        eos_token_id=tokenizer.eos_token_id,
        max_new_tokens=512,
        do_sample=False if temperature == 0 else True,
        use_cache=True,
        temperature=temperature,
        top_p=top_p,
    )
    
    answer = tokenizer.decode(outputs[0].cpu().tolist(), skip_special_tokens=True)
    return answer

from util_for_os import ose,osj

# date = '0808'

# dates = ['0818' , '0819']
# dates = ['0820']
# dates = ['0821','0822']
# dates = ['0825']
# dates = ['0826']
# dates = ['0827','0828']
# dates = ['0829']
# dates = ['0901']
# dates = ['0902']
# dates = ['0903','0904']
# dates = ['0905']
# dates = ['0908']
# dates = ['0925','0926','0928']

# categories = ['coat','leather','sweater']
# categories = ['coat','sweater']
# categories = ['dress','slipdress','feather']
# categories = ['dress',]


dates = ['1106']
categories = ['jacket','hoodie','formalatt','weddress']


# dates = ['0917']
# categories = ['slipdress']

for date in dates:
    for category in categories:
        dir = f'/mnt/nas/datasets/diction/{category}{date}_img_clo_diff'
        assert ose( dir ), dir
        assert ose( osj( dir , 'names.txt' ) )
        
for date in dates:
    for category in categories:
        dir = f'/mnt/nas/datasets/diction/{category}{date}_img_clo_diff'
# dir = f'/mnt/nas/datasets/diction/coat{date}_img_clo_diff'
# dir = f'/mnt/nas/datasets/diction/leather{date}_img_clo_diff'
# dir = f'/mnt/nas/datasets/diction/sweater{date}_img_clo_diff'

        save_clothing_dir = f'{dir}_pattern'
        save_human_dir = f'{dir}_patternless'
        import os,shutil,cv2
        # filenames = os.listdir(dir)
        if os.path.exists(save_clothing_dir):shutil.rmtree(save_clothing_dir)
        os.makedirs(save_clothing_dir)
        if os.path.exists(save_human_dir):shutil.rmtree(save_human_dir)
        os.makedirs(save_human_dir)

        Q1 = '如何描述这件衣服的图案'  # no
        # Q2 = '是否是一件完整的衣服'   # yes
        # Q3 = 'Is the background white? Just answer yes or no.'   # yes
        # Q4 = '是否是鞋子' # no
        # Q5 = '是否是衣服' # yes

        keywords = ['重复','规则','网格',
                    '千鸟格']
        negative_keywords = ['绗缝','皱褶','褶皱','不规则']

        # check_no = lambda ans : 'no' in ans.lower() or '否' in ans or '不是' in ans
        check_yes = lambda ans : any( [ k in ans for k in keywords] ) and \
                                all( [ k not in ans for k in negative_keywords ] )

        # count = 0
        with open(  osj( dir , 'names.txt' ) ,encoding='utf-8') as f:
            names = f.readlines() 
        import pdb
        from tqdm import tqdm
        # pdb.set_trace()
        # for entry in os.scandir( dir ):
        get_answer = lambda filepath,q : multimodal_understanding( cv2.imread( filepath ) , q )

        will_save = []
        save_json_path = osj( dir , 'tmp_save_pattern.json' ) 

        from util_for_os import ose
        import json
        # pdb.set_trace()
        if ose(save_json_path):
            with open( save_json_path , encoding='utf-8' ) as f:
                will_save = json.load(f)
                
            for data in tqdm(will_save):
                filepath = data['filepath']
                ans = data['ans']
                
                filename = os.path.basename(filepath)
                
                if check_yes( ans ) :
                    shutil.copy2( osj(dir , filename) , osj(save_clothing_dir,filename) )
                else:
                    shutil.copy2( osj(dir , filename) , osj(save_human_dir,filename) )

        else:
            model_path = "/mnt/nas/zhangshu/Janus-Pro-7B"
            config = AutoConfig.from_pretrained(model_path) if config is None else config
            language_config = config.language_config if language_config is None else language_config
            language_config._attn_implementation = 'eager'
            vl_gpt = AutoModelForCausalLM.from_pretrained(model_path,
                                                        language_config=language_config,
                                                        trust_remote_code=True) if vl_gpt is None else vl_gpt
            if torch.cuda.is_available():
                vl_gpt = vl_gpt.to(torch.bfloat16).cuda()
            else:
                vl_gpt = vl_gpt.to(torch.float16)

            vl_chat_processor = VLChatProcessor.from_pretrained(model_path) if vl_chat_processor is None else vl_chat_processor
            tokenizer = vl_chat_processor.tokenizer
            cuda_device = 'cuda' if torch.cuda.is_available() else 'cpu'
            
            for filename in tqdm(names):
                # if entry.is_file() and entry.name.endswith('.jpg'):
                filename = filename.strip()
                if filename.endswith('.jpg'):
                    # filename = entry.name
                    pass
                else:continue

                # count += 1
                # print('\rprocess num : ',count,end='',flush=True)    
                
                filepath = os.path.join( dir , filename)

                ans = get_answer( filepath , Q1 )
                
                will_save.append({
                    'filepath':filepath,
                    'ans':ans,
                })
                
                # pdb.set_trace()
                
                if check_yes( ans ) :
                    shutil.copy2( osj(dir , filename) , osj(save_clothing_dir,filename) )
                else:
                    shutil.copy2( osj(dir , filename) , osj(save_human_dir,filename) )
            
            import json
            with open( save_json_path ,'w',encoding='utf-8') as f:
                json.dump( will_save , f , ensure_ascii=False , indent=2 )
'''
/mnt/nas/datasets/diction/images/               筛选后的图(最终存放位置)
/mnt/nas/datasets/diction/ZipArchive            不筛选的zip
/mnt/nas/datasets/diction/ZipArchive07XX        解压后的图片
/mnt/nas/datasets/diction/ZipArchive07XX_clo    Janus筛选后的clo图 (包含重复图,可能出现几张不完整或者背景不干净,错误概率比例小于2%)
/mnt/nas/datasets/diction/ZipArchive07XX_clo_sim     Janus筛选后的clo图 (包含重复图 clo)
/mnt/nas/datasets/diction/ZipArchive07XX_clo_diff    Janus筛选后的clo图 (包含不重复图 clo,可能出现几张不完整或者背景不干净,错误概率比例小于2%)
/mnt/nas/datasets/diction/ZipArchive07XX_human      Janus筛选后的非clo图 (人 / 不完整)

date=0902
bash restore.sh /mnt/nas/datasets/diction/coat"$date"
bash restore.sh /mnt/nas/datasets/diction/sweater"$date"
bash restore.sh /mnt/nas/datasets/diction/leather"$date"
bash restore.sh /mnt/nas/datasets/diction/jean"$date"

date=0902
bash restore.sh /mnt/nas/datasets/diction/coat"$date"_img
bash restore.sh /mnt/nas/datasets/diction/sweater"$date"_img
bash restore.sh /mnt/nas/datasets/diction/leather"$date"_img
bash restore.sh /mnt/nas/datasets/diction/jean"$date"_img

date=1103
bash restore.sh /mnt/nas/datasets/diction/coat"$date"_img_clo
bash restore.sh /mnt/nas/datasets/diction/sweater"$date"_img_clo
bash restore.sh /mnt/nas/datasets/diction/leather"$date"_img_clo
bash restore.sh /mnt/nas/datasets/diction/jean"$date"_img_clo

date=1106
bash restore.sh /mnt/nas/datasets/diction/jacket"$date"_img_clo_diff
bash restore.sh /mnt/nas/datasets/diction/hoodie"$date"_img_clo_diff
bash restore.sh /mnt/nas/datasets/diction/formalatt"$date"_img_clo_diff
bash restore.sh /mnt/nas/datasets/diction/weddress"$date"_img_clo_diff


'''